import sys
import mss.tools
import numpy as np
import cv2
import signal
from ultralytics.yolo.engine.results import Results
from ultralytics import YOLO
from typing import List, Tuple, Union
import win32api
import win32con

# 定义配置
precision = 0.7  # 识别精度
max_det = 1  # 单次识别目标数量
verbose = False  # 是否显示日志
SCREEN_W = 1920  # 屏幕长
SCREEN_H = 1080  # 屏幕高
SCREENSHOT_W = 640  # 截图区域长
SCREENSHOT_H = 640  # 截图区域高
SCREEN_CX = SCREEN_W // 2  # 屏幕中心x
SCREEN_CY = SCREEN_H // 2  # 屏幕中心y
LEFT = SCREEN_CX - SCREENSHOT_W // 2  # 检测框左上角x
TOP = SCREEN_CY - SCREENSHOT_H // 2  # 检测框左上角y

# 程序初始化
model = YOLO('CF-Best.pt')
sct = mss.mss()


# 获取屏幕截图
def get_screenshot() -> np.array:
    return np.array(sct.grab({"top": TOP, "left": LEFT, "width": SCREENSHOT_W, "height": SCREENSHOT_H}))


# 目标检测
def object_detection(img: np.array) -> List[Results]:
    return model.predict(cv2.cvtColor(img, cv2.COLOR_RGBA2RGB), conf=precision, imgsz=[SCREENSHOT_W, SCREENSHOT_H], max_det=max_det, verbose=verbose, save=False)


# 根据目标识别结果获取射击点
def object_center(objects: List[Results]) -> Union[None, Tuple[int, int]]:
    # 获取识别目标后的坐标
    obj_points = objects[0].boxes.xyxy

    # 如果识别目标成功，则计算射击点
    if len(obj_points) != 0:
        left, top, right, bottom = get_coordinate(obj_points)

        x = int(LEFT + ((left + right) // 2) - SCREEN_CX)
        y = int(TOP + ((top + bottom) // 2) - SCREEN_CY)
        print(x, y)
        return x, y
    else:
        return None


# 获取坐标点 left,top,right,bottom
def get_coordinate(xy_xy):
    return xy_xy[0][0], xy_xy[0][1], xy_xy[0][2], xy_xy[0][3]


# 移动鼠标
def move_mouse(x: int, y: int):
    win32api.mouse_event(win32con.MOUSEEVENTF_MOVE, int(x), int(y), 0, 0)


# 主进程退出
def main_exit(event, frame):
    sys.exit(0)


# 程序启动
if __name__ == '__main__':
    signal.signal(signal.SIGINT, main_exit)
    signal.signal(signal.SIGTERM, main_exit)

    while True:
        # 获取区域截图
        image = get_screenshot()

        # 将截图进行目标识别
        res = object_detection(image)

        # 获取射击点
        points = object_center(res)
        if points is None:
            continue

        # 移动鼠标到目标点
        move_mouse(points[0], points[1])
